"""
自定义collate函数，用于处理数据类型不匹配问题
"""

import torch
import numpy as np
from torch.utils.data._utils.collate import default_collate

def custom_collate_fn(batch):
    """
    自定义collate函数，确保所有张量类型一致
    """
    if len(batch) == 0:
        return batch
    
    # 获取第一个样本作为参考
    first_sample = batch[0]
    
    # 处理每个键
    collated_batch = {}
    for key in first_sample.keys():
        try:
            # 收集所有样本的该键值
            values = [sample[key] for sample in batch]
            
            # 检查数据类型一致性
            if isinstance(values[0], torch.Tensor):
                # 确保所有张量类型一致
                target_dtype = values[0].dtype
                for i, val in enumerate(values):
                    if val.dtype != target_dtype:
                        if target_dtype == torch.bool:
                            values[i] = val.bool()
                        elif target_dtype == torch.float32:
                            values[i] = val.float()
                        elif target_dtype == torch.int16:
                            values[i] = val.short()
                        elif target_dtype == torch.int32:
                            values[i] = val.int()
                        elif target_dtype == torch.int64:
                            values[i] = val.long()
                        else:
                            values[i] = val.to(target_dtype)
                
                # 使用torch.stack进行拼接
                try:
                    collated_batch[key] = torch.stack(values, dim=0)
                except RuntimeError as e:
                    if "input types can't be cast" in str(e):
                        # 如果类型转换失败，尝试统一转换为float32
                        print(f"警告: 键 {key} 类型转换失败，统一转换为float32")
                        float_values = [val.float() for val in values]
                        collated_batch[key] = torch.stack(float_values, dim=0)
                    else:
                        raise e
            else:
                # 非张量数据使用默认collate
                collated_batch[key] = default_collate(values)
                
        except Exception as e:
            print(f"处理键 {key} 时出错: {e}")
            print(f"数据类型: {type(values[0]) if values else 'None'}")
            if values and isinstance(values[0], torch.Tensor):
                print(f"张量形状: {values[0].shape}")
                print(f"张量类型: {values[0].dtype}")
            # 尝试使用默认collate作为后备
            try:
                collated_batch[key] = default_collate(values)
            except Exception as e2:
                print(f"默认collate也失败: {e2}")
                # 如果都失败，跳过这个键
                continue
    
    return collated_batch

def safe_collate_fn(batch):
    """
    安全的collate函数，处理各种数据类型问题
    """
    if len(batch) == 0:
        return batch
    
    # 预处理：确保所有张量都是相同的数据类型
    processed_batch = []
    for sample in batch:
        processed_sample = {}
        for key, value in sample.items():
            if isinstance(value, torch.Tensor):
                # 根据键名决定数据类型
                if 'center_image' in key:
                    processed_sample[key] = value.float()
                elif 'image' in key:
                    processed_sample[key] = value.float()
                elif 'flow' in key:
                    processed_sample[key] = value.float()
                elif 'instance' in key or 'label' in key:
                    processed_sample[key] = value.short()
                else:
                    processed_sample[key] = value.float()  # 默认使用float
            else:
                processed_sample[key] = value
        processed_batch.append(processed_sample)
    
    # 使用自定义collate函数
    return custom_collate_fn(processed_batch) 